import cv2
import numpy as np

from PIL import Image
from transparent_background import Remover

# Load model
remover = Remover(mode='fast',device='cuda:0',) # nightly release checkpoint

# Usage for video
cap = cv2.VideoCapture(r'G:/clip/UID62841292067_火只木南/video/2023-05-13 16.23.59-火只木南-0717 y@T.lC 在座的各位 敢不敢实名观看？ # 极乐净土花式挑战 # 盛典与慧业 # 八重神子/2023-05-13 16.23.59-火只木南-0717 y@T.lC 在座的各位 敢不敢实名观看？ # 极乐净土花式挑战 # 盛典与慧业 # 八重神子-Scene-001.mp4')  # video reader for input
fps = cap.get(cv2.CAP_PROP_FPS)

writer = None

while cap.isOpened():
    ret, frame = cap.read()  # read video

    if ret is False:
        break

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    img = Image.fromarray(frame).convert('RGB')

    if writer is None:
        writer = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps,
                                 img.size)  # video writer for output

    out = remover.process(img, type='blur')  # same as image, except for 'rgba' which is not for video.
    writer.write(cv2.cvtColor(np.array(out), cv2.COLOR_BGR2RGB))

cap.release()
writer.release()